package com.fanshuai.metrics.config;

import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.MetricSet;
import com.codahale.metrics.jvm.*;
import io.prometheus.client.dropwizard.DropwizardExports;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;

import javax.management.MBeanServer;
import java.lang.management.ManagementFactory;
import java.util.Map;

@Configuration
@EnableAutoConfiguration(
        exclude = {DataSourceAutoConfiguration.class}
)
@ComponentScan("com.fanshuai.metrics")
public class MetricsPrometheusConfig {

    @Bean
    public DropwizardExports dropwizardExports() {
        MetricRegistry registry = new MetricRegistry();
        MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();

        registerAll("gc", new GarbageCollectorMetricSet(), registry);
        registerAll("buffers", new BufferPoolMetricSet(mBeanServer), registry);
        registerAll("memory", new MemoryUsageGaugeSet(), registry);
        registerAll("threads", new ThreadStatesGaugeSet(), registry);
        registerAll("classLoading", new ClassLoadingGaugeSet(), registry);
        registry.register("file.descriptor.ratio", new FileDescriptorRatioGauge());

        return new DropwizardExports(registry);
    }

    private static void registerAll(String prefix, MetricSet metricSet, MetricRegistry registry) {
        Map<String, Metric> metricMap = metricSet.getMetrics();
        for (String metricName : metricMap.keySet()) {
            Metric metric = metricMap.get(metricName);

            if (metric instanceof MetricSet) {
                registerAll(prefix + "." + metricName, (MetricSet) metric, registry);
            } else {
                registry.register(prefix + "." + metricName, metric);
            }
        }
    }
}
